#include "../include/model.h"
#include "../include/process.h"
#include "../include/rknn_api.h"
#include "../include/readfile.h"
#include <cstdio>
#include <set>

const int height = 640;
const int width = 640;
const int channel = 3;
const char* model = "./model.rknn";
std::string coco_class[80] = {"person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
           "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant",
           "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
           "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ",
           "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa",
           "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop","mouse","remote ","keyboard ","cell phone","microwave ",
           "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush"};

int model_init(rknn_ctx* rknn)
{
    int model_len = 0;
    rknn->out_data = NULL;

    // 模型加载并返回模型二进制数据
    model_len = read_data_from_file(model, &(rknn->out_data));
    if (rknn->out_data == NULL)
    {
        printf("load_model fail!\n");
        return -1;
    }

    // rknn 初始化
    rknn->rknn_ctx = 0;
    int ret = rknn_init(&(rknn->rknn_ctx), rknn->out_data, model_len, 0);
    if (ret < 0)
    {
        printf("rknn_init fail! ret=%d\n", ret);
        return -1;
    }

    // 得到input output的num
    ret = rknn_query(rknn->rknn_ctx, RKNN_QUERY_IN_OUT_NUM, &(rknn->io_num), sizeof(rknn->io_num));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", rknn->io_num.n_input, rknn->io_num.n_output);

    // Get Model Input Info
    rknn->input_attrs = (rknn_tensor_attr*)malloc(sizeof(rknn_tensor_attr)*(rknn->io_num.n_input));
    for (int i = 0; i < rknn->io_num.n_input; i++)
    {
        rknn->input_attrs[i].index = i;
        ret = rknn_query(rknn->rknn_ctx, RKNN_QUERY_INPUT_ATTR, &(rknn->input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            return -1;
        }
    }

    // Get Model Output Info
    rknn->output_attrs = (rknn_tensor_attr*)malloc(sizeof(rknn_tensor_attr)*(rknn->io_num.n_output));
    for (int i = 0; i < rknn->io_num.n_output; i++)
    {
        rknn->output_attrs[i].index = i;
        ret = rknn_query(rknn->rknn_ctx, RKNN_QUERY_OUTPUT_ATTR, &(rknn->output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            return -1;
        }
    }

    return 0;
}

int model_inference(rknn_ctx* rknn, void* data, std::vector<float> &filterBoxes, 
                    std::vector<float> &objProbs, std::vector<int> &classId, std::vector<int> &indexArray)
{
    // set input
    rknn_input inputs[rknn->io_num.n_input];
    memset(inputs, 0, sizeof(inputs));
    for(int i = 0; i < rknn->io_num.n_input; i++)
    {
        inputs[i].index = i;
        inputs[i].type = RKNN_TENSOR_UINT8;
        inputs[i].fmt = RKNN_TENSOR_NHWC;
        inputs[i].size = channel*width*height;
        inputs[i].buf = data;
    }
    int ret = rknn_inputs_set(rknn->rknn_ctx, rknn->io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("rknn_input_set fail! ret=%d\n", ret);
        return -1;
    }

    // inference
    ret = rknn_run(rknn->rknn_ctx, nullptr);
    if (ret < 0)
    {
        printf("rknn_run fail! ret=%d\n", ret);
        return -1;
    }

    // get output
    rknn_output outputs[rknn->io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < rknn->io_num.n_output; i++)
    {
        outputs[i].index = i;
        outputs[i].want_float = false;
    }
    ret = rknn_outputs_get(rknn->rknn_ctx, rknn->io_num.n_output, outputs, NULL);
    if (ret < 0)
    {
        printf("rknn_outputs_get fail! ret=%d\n", ret);
        return -1;
    }

    // 数据后处理
    int validCount = 0;
    int stride = 0;
    int grid_h = 0;
    int grid_w = 0;
    int dfl_len = rknn->output_attrs[0].dims[2] / 4;
    int output_per_branch = rknn->io_num.n_output / 3;
    for (int i = 0; i < output_per_branch; i++)
    {
        void *score_sum = nullptr;
        int32_t score_sum_zp = 0;
        float score_sum_scale = 1.0;
        if (output_per_branch == 3){
            score_sum = outputs[i*output_per_branch + 2].buf;
            score_sum_zp = rknn->output_attrs[i*output_per_branch + 2].zp;
            score_sum_scale = rknn->output_attrs[i*output_per_branch + 2].scale;
        }
        int box_idx = i*output_per_branch;
        int score_idx = i*output_per_branch + 1;
        grid_h = rknn->output_attrs[box_idx].dims[1];
        grid_w = rknn->output_attrs[box_idx].dims[0];
        stride = height / grid_h;
        validCount += process_u8((uint8_t *)outputs[box_idx].buf, rknn->output_attrs[box_idx].zp, rknn->output_attrs[box_idx].scale,
                                (uint8_t *)outputs[score_idx].buf, rknn->output_attrs[score_idx].zp, rknn->output_attrs[score_idx].scale,
                                (uint8_t *)score_sum, score_sum_zp, score_sum_scale,
                                grid_h, grid_w, stride, dfl_len,
                                filterBoxes, objProbs, classId, 0.32);
    }

    if(validCount == 0)
        return 0;

    for (int i = 0; i < validCount; ++i)
    {
        indexArray.push_back(i);
    }
    std::set<int> class_set(std::begin(classId), std::end(classId));
    for (auto c : class_set)
    {
        nms(validCount, filterBoxes, classId, indexArray, c, 0.45);
    }

    return validCount;
}